import pandas as pd
import matplotlib.pyplot as plt
from sqlalchemy import create_engine
from config import MYSQL_CONFIG

# 设置中文字体和负号显示正常
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False


# 连接数据库并获取数据
def fetch_data(table_name):
    db_url = f"mysql+pymysql://{MYSQL_CONFIG['user']}:{MYSQL_CONFIG['password']}@{MYSQL_CONFIG['host']}:{MYSQL_CONFIG['port']}/{MYSQL_CONFIG['database']}"
    engine = create_engine(db_url)
    query = f"SELECT * FROM `{table_name}`"
    df = pd.read_sql(query, engine)
    return df


# 在一个窗口中绘制三张柱状图：合格率、司机数量、燃油费用
def plot_three_charts_in_one_window(df):
    if df.empty:
        print("数据为空，无法绘图")
        return

    # 过滤掉 logistics_company 为 None 的行
    df = df[df['logistics_company'].notna()]

    # 获取唯一的物流公司名称作为 X 轴
    logistics_companies = df['logistics_company'].unique()

    # 准备子图画布
    fig, axes = plt.subplots(1, 3, figsize=(18, 6), sharex=True)

    # --- 图表 1：产品合格率 ---
    if 'all_acount' in df.columns and 'damage_count' in df.columns:
        df['合格率'] = (df['all_acount'] - df['damage_count']) / df['all_acount'].replace(0, pd.NA).fillna(0)
        product_quality = df.groupby('logistics_company')['合格率'].mean().reindex(logistics_companies).fillna(0)

        product_quality.plot(kind='bar', ax=axes[0], color='skyblue')
        axes[0].set_title('各物流公司产品合格率')
        axes[0].set_ylabel('合格率')
        axes[0].tick_params(axis='x', rotation=45)
    else:
        axes[0].axis('off')  # 如果字段缺失，则隐藏该图

    # --- 图表 2：司机数量 ---
    if 'logistics_company' in df.columns and 'driver_id' in df.columns:
        driver_count = df.groupby('logistics_company')['driver_id'].nunique().reindex(logistics_companies).fillna(0)

        driver_count.plot(kind='bar', ax=axes[1], color='lightgreen')
        axes[1].set_title('各物流公司司机数量（去重统计）')
        axes[1].set_ylabel('司机人数')
        axes[1].tick_params(axis='x', rotation=45)
    else:
        axes[1].axis('off')

    # --- 图表 3：燃油费用 ---
    if 'logistics_company' in df.columns and 'fuel_consumed_per_km' in df.columns:
        df['fuel_consumed_per_km'] = pd.to_numeric(df['fuel_consumed_per_km'], errors='coerce').fillna(0)
        fuel_cost = df.groupby('logistics_company')['fuel_consumed_per_km'].sum().reindex(logistics_companies).fillna(0)

        fuel_cost.plot(kind='bar', ax=axes[2], color='salmon')
        axes[2].set_title('各物流公司燃油费用总和')
        axes[2].set_ylabel('燃油消耗总量')
        axes[2].tick_params(axis='x', rotation=45)
    else:
        axes[2].axis('off')

    plt.tight_layout()
    plt.show()


if __name__ == "__main__":
    table_name = "transport_task_final"
    df = fetch_data(table_name)

    if df.empty:
        print("数据为空，请检查数据库是否包含有效数据")
    else:
        print("成功读取数据！以下是可用列名：")
        print(df.columns.tolist())

        plot_three_charts_in_one_window(df)